import cv2
import matplotlib.pyplot as plt
import numpy as np
import ezdxf
doc = ezdxf.new(dxfversion='R2010')  # 创建新 DXF 文档
msp = doc.modelspace()  # 获取模型空间
# 读取图像
image = cv2.imread('./data/aa.png')

# 原始图像副本
original_image = image.copy()

# 转换为灰度图像
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
# 调整后的参数
edges = cv2.Canny(gray, 80, 200, apertureSize=3)
lines = cv2.HoughLines(edges, rho=1, theta=np.pi / 180, threshold=100)
# 创建一个空白图像用于绘制线段（与原始图像大小相同）
line_image = np.zeros_like(original_image)

if lines is not None:
    for line in lines:
        rho, theta = line[0]
        a = np.cos(theta)
        b = np.sin(theta)
        x0 = a * rho
        y0 = b * rho
        # 计算线段的两个点（长度为1000像素）
        x1 = int(x0 + 1000 * (-b))
        y1 = int(y0 + 1000 * (a))
        x2 = int(x0 - 1000 * (-b))
        y2 = int(y0 - 1000 * (a))

        # 绘制直线
        color = tuple(np.random.randint(150, 255, 3).tolist())
        cv2.line(line_image, (x1, y1), (x2, y2), color, 2)

        # 添加直线到 DXF
        msp.add_line((x1, y1), (x2, y2), dxfattribs={'color': 7})
        
plt.figure(figsize=(18, 6))
doc.saveas('output/detected_lines.dxf')  # 保存 DXF 文件
# 原始图像
plt.subplot(1, 3, 1)
plt.imshow(cv2.cvtColor(original_image, cv2.COLOR_BGR2RGB))
plt.title("Original Image")
plt.axis('off')

# 边缘检测图像
plt.subplot(1, 3, 2)
plt.imshow(edges, cmap='gray')  # 显示 Canny 边缘图像
plt.title("Canny Edges")
plt.axis('off')

# 线段检测结果（无原图背景）
plt.subplot(1, 3, 3)
plt.imshow(cv2.cvtColor(line_image, cv2.COLOR_BGR2RGB))
plt.title("Detected Lines (No Background)")
plt.axis('off')

plt.tight_layout()
plt.show()